
'''
DATASET_NAME
├─ camera[_TYPE].json
├─ dataset_info.json
├─ test_targets_bop19.json
├─ models[_MODELTYPE][_eval]
│  ├─ models_info.json
│  ├─ obj_OBJ_ID.ply
├─ train|val|test[_TYPE]
│  ├─ SCENE_ID|OBJ_ID
│  │  ├─ scene_camera.json
│  │  ├─ scene_gt.json
│  │  ├─ scene_gt_info.json
│  │  ├─ depth
│  │  ├─ mask
│  │  ├─ mask_visib
│  │  ├─ rgb|gray
'''

# 遍历所有的txt文件 并保存为npy文件
import os
import numpy as np
import json
from PIL import Image
from glob import glob

try:
    import ruamel_yaml as yaml
except ModuleNotFoundError:
    import ruamel.yaml as yaml


from bbox.box_options import box_cxcywh_to_xyxy

#读取一个json文件

def read_json(json_path):
    with open(json_path, 'r') as f:
        data = json.load(f)
    return data

dataset=read_json(r'E:\pose\label\project-2-at-2024-03-16-14-00-7af4688b\result.json')
dataset_length= len(dataset['images'])


def read_txt(txt_path):
    data=np.loadtxt(txt_path)
    return data

def process_pose(pose):
    #转为3*4矩阵
    pose = np.array(pose).reshape(3, 4)
    #转为4*4矩阵
    pose = np.concatenate([pose, np.array([[0, 0, 0, 1]])], axis=0)
    return pose

def save_dict_to_json(dicts:dict, json_path:str):
    data= json.dumps(dicts, indent=4)
    with open(json_path, 'w') as f:
        json.dump(dicts,f)


def create_gt_json(root=r'E:\pose\datasets\linemod\LINEMOD'):
    all_poses = {}
    for idx in range(dataset_length):
        pose_obj_path = os.path.join(root, r"objects\target\pose\mtrix","{:06d}.txt".format(idx)) #E:\pose\datasets\linemod\LINEMOD\objects\target\pose\mtrix
        pose = read_txt(pose_obj_path)
        print(pose)
        pose = process_pose(pose)
        all_poses[idx] = {'cam_R_w2c': pose[:3, :3], 'cam_t_w2c': pose[:3, 3], 'obj_id': 1}
    save_info(r'train_gt.json',all_poses)

def save_info(path, info, save_all=False):
    if not save_all:
        for im_id in sorted(info.keys()):
            im_info = info[im_id]
            if 'cam_K' in im_info.keys():
                im_info['cam_K'] = im_info['cam_K'].flatten().tolist()
            if 'cam_R_w2c' in im_info.keys():
                im_info['cam_R_w2c'] = im_info['cam_R_w2c'].flatten().tolist()
            if 'cam_t_w2c' in im_info.keys():
                im_info['cam_t_w2c'] = im_info['cam_t_w2c'].flatten().tolist()
            if 'obj_id' in im_info.keys():
                im_info['obj_id'] = int(im_info['obj_id'])
        with open(path, 'w') as f:
            json.dump(info, f)
    else:
        for im_id in sorted(info.keys()):
            im_info = info[im_id]
            for key in im_info.keys():
                if key in ['obj_bb', 'cam_R_m2c', 'cam_t_m2c', 'cam_K']:
                    im_info[key] = np.asarray(im_info[key]).flatten().tolist()
        with open(path, 'w') as f:
            json.dump(info, f, Dumper=yaml.CDumper, width=10000)



def load_json(json_file):
    with open(json_file, 'r') as f:
        data = json.load(f)
    return data



def create_gt_info_json(source_path, target_path='train_gt_info.json'):
    data = load_json(source_path)
    all_info = {}
    for idx in range(dataset_length):
        im_info = data['annotations'][idx]
        bbox_coco=im_info['bbox']
        bbox_xyxy= [bbox_coco[0],bbox_coco[1],bbox_coco[0]+bbox_coco[2],bbox_coco[1]+bbox_coco[3]]
        bbox_xyxy = [int(i) for i in bbox_xyxy]
        all_info[idx] = {'bbox_obj':bbox_xyxy} # 取整数
        
    save_info(target_path, all_info)


create_gt_json()

# create_gt_info_json(r'E:\pose\label\project-2-at-2024-03-16-14-00-7af4688b\result.json')





